# -*-coding=utf-8-*-

"""
5参数的势能面
"""

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np


class SimpleSurface():
    def __init__(self):
        self.point = []
        self.fig = None

    def get_value(self, vector):
        """
        计算函数
        """
        return np.sum(np.sin(vector))

    def get_force(self, vector):
        """
        计算力
        """
        return -np.cos(vector)

    def show_surface(self):
        x_array = np.arange(-8, 8, 0.5)
        y_array = np.arange(-8, 8, 0.5)
        x, y = np.meshgrid(x_array, y_array)
        z = [self.get_value(x_, y_) for x_, y_ in zip(x, y)]
        # self.ax.plot_surface(x, y, z, rstride=1, cstride=1, cmap=plt.get_cmap('rainbow'))
        self.ax.scatter(x, y, z, c='r')
        # ax.contourf(x, y, z, zdir='z', offset=2)
        # plt.show()

    def show_point(self, x, y, z):
        """
        显示点
        """
        self.fig = plt.figure()
        # ax = fig.add_subplot(111, projection='3d')
        self.ax = Axes3D(self.fig)
        self.ax.scatter(x, y, z)
        self.ax.set_zlim(-4, 4)

    def show(self):
        plt.show()


if __name__ == "__main__":
    ss = SimpleSurface()
    ss.show_surface()
    x = [1, 1]
    y = [1, 2]
    z = [3, 3]
    ss.show_point(x, y, z)
    ss.show()
